#include "base.h"
#include "aux.h"
#include "sfhash.h"
#include "prefix.hpp"
#include "trie.hpp"
#include "asnprefix.hpp"
#include "cflow.hpp"

using namespace std;

const string version = "0.07";

//#define DEBUG_ENABL6

#ifdef DEBUG_ENABLE
#   define DEBUGPREFIX "Trie"
#   define debugsn(x) { cout << "(" << DEBUGPREFIX << ") " << x << endl; }
#else
#   define debugsn(x...)
#endif
//#define debug(x...) printf("(%s) ",DEBUGPREFIX); printf(x);

/* CITM, AS - prefix */
class CITMValue {
        uint64_t flows;
        uint64_t octets;
        uint64_t packets;
    public:
        CITMValue() { flows = octets = packets = 0; }
        CITMValue( uint64_t _flows, uint64_t _octets, uint64_t _packets ) {
            flows = _flows; octets = _octets; packets = _packets;
        }

        uint64_t getFlows()   { return flows;   }
        uint64_t getOctets()  { return octets;  }
        uint64_t getPackets() { return packets; }

        CITMValue operator=( CITMValue rhs ) {
            if ( this == &rhs ) return *this;
            flows   = rhs.getFlows();
            octets  = rhs.getOctets();
            packets = rhs.getPackets();
            return *this;
        }

        CITMValue operator+( CITMValue rhs ) {
            CITMValue result( this->flows   + rhs.getFlows(),
                             this->octets  + rhs.getOctets(),
                             this->packets + rhs.getPackets() );
            return result;
        }

        CITMValue operator+=( CITMValue rhs ) {
            this->flows   += rhs.getFlows();
            this->octets  += rhs.getOctets();
            this->packets += rhs.getPackets();
            return *this;
        }

        string str( string delimiter = " ") {
            ostringstream oss;
            oss << flows << delimiter << octets << delimiter << packets;
            return oss.str();
        }
};

//class CMKey {
//    public:
//        string virtual str()=0;
//	    class LessThan : public std::binary_function<CMKey, CMKey, bool>{};
//};

//  template <class CITMKey, class CITMKeyLess >
//  class CITM {
//  
//  
//  
//      private:
//  
//  
//          /* map to store AS-prefix */
//          /* key*/
//          map<CITMKey, CITMValue, CITMKeyLess> m;
//  
//      public:
//          CITM(){ 
//  //            m = new map<CITMKey,CITMValue>; // TODO add we<k compare reference
//          };
//          ~CITM() {
//              // TODO add destructor for all the entries
//          }
//  
//          CITMValue add( CITMKey k , CITMValue v) {
//              /* m[k] wil create an entry if this does not exist... */
//              m[k] += v;
//          }
//  
//          CITMValue get( CITMKey k ) {
//              return m[k];
//          }
//  
//          /* iterator */
//          map<CITMKey, CITMValue, CITMKeyLess>::iterator begin() {
//              return m.begin();
//          }
//          
//  
//  
//  
//  };

/** -- **/


#define KEYFORMAT_ASNPREFIX 1
#define KEYFORMAT_ASNASN    2

#define INTYPE_NETFLOW  1
#define INTYPE_SMARTBIN 2

/* options */
struct {
    string file_in; 
    int format;
    list<string> files_mapping;
    int keyFormat;
    int inType;
} globalArgs = { "-", 1, list<string>(), KEYFORMAT_ASNPREFIX, INTYPE_NETFLOW} ;

const char * shortOpts = "i:m:f:k:I:";

const struct option longOpts[] = {
    {"in", required_argument, NULL, 'i' },
    {"mapping", required_argument, NULL, 'm' },
    {"format", required_argument, NULL, 'f' },
    {"aggregation", required_argument, NULL, 'k'},
    {"in-type", required_argument, NULL, 'I'},
    {"help", no_argument, NULL, 'h'},
    { NULL, no_argument, NULL, 0}
};

void display_usage() {  
    cout << "version: " << version << endl;
    cout << "--help -h          help\n";
    cout << "--in -i            input file, can occur multiple times\n";
    cout << "--mapping -m       file with prefix-asn mapping (can multiple)\n";
    cout << "                   mandatory for 'smart' file format\n";
    cout << "--format -f        1 : as_src as_dst flows octets packets\n";
    cout << "                   2 : as_src as_dst prefix_dst flows octets packets\n"; 
    cout << "--aggregation -a   aa : aggregate by src. AS, dst. AS\n";
    cout << "                   ap : aggregate by src. AS, dst. prefix\n";
    cout << "--in-type -I       nf netflow : NetFlow v5\n";
    cout << "                   s smart : smartxac binary format\n";
}

string format_info[] = {
    "dummy",
    "src_as dst_prefix flows octets packets",
    "src_as dst_as dst_prefix flows octets packets"
};

void display_information() {
    list<string>::iterator it;
    time_t rawtime;

    time ( &rawtime );

    cout << "# txas++ , version " << version << endl;
    cout << "# time: " << ctime (&rawtime);
    cout << "# timestamp: " << rawtime << endl;
    cout << "# format: " << globalArgs.format << endl;
    switch ( globalArgs.keyFormat ) {
        case KEYFORMAT_ASNASN: cout << "# aggregation: as-as" << endl; break;
        case KEYFORMAT_ASNPREFIX: cout << "# aggregation: as-prefix" << endl; break;
    }

    if ( globalArgs.files_mapping.empty() ) {
        cout << "# mapping: (none)\n";
    }
    else {
        for( it = globalArgs.files_mapping.begin(); it != globalArgs.files_mapping.end(); it++) {
            cout << "# mapping: " << *it << endl;
        }
    }

    cout << "# in: " << globalArgs.file_in << endl;

    switch( globalArgs.inType ) {
        case INTYPE_SMARTBIN: cout << "# input type: smart bin\n"; break;
        case INTYPE_NETFLOW:  cout << "# input type: netflow\n"; break;
    }

    cout << "#" << endl;
    cout << "" << format_info[globalArgs.format] << endl;
}

/* main */

int main( int argc, char ** argv) {
#if 0 
    CTrieTree * t;
     CTrieNode * n;
    printf( "%08x\n", 0xabcdef12 & ~(0xFFFFFFFF >> 8));
    printf( ">>> %08x\n", mask32( 0xabcdef12 , 32));
     t->insertNode( 100, new Prefix(0x00000000, 0) );
     t->insertNode( 200, new Prefix(0xffffffff, 32) );
     t->insertNode( 300, new Prefix(0xffffffff, 8 ) );
 //    t->insertNode( 100, new Prefix(0x00010000, 32) );
 
     cout << t->lpmatch(0xffffffff)->str() << "\n";
     cout << t->lpmatch(0xffffffee)->str() << "\n";
#endif 

# if 0
    string fname_mapping = "../data/prefixes/prefixes.2011-06-05-1800.txt";

    CTrieTree * t;
    t = CTrieTree::newFromFile(fname_mapping, CTrieTree::FORMATTXT);
    t->addSpecial(); // TODO force overwriting nodes ?
//    cout << t->lpmatch( bytes_4_to_uint32( 147,83,35,173  )  )->str() << "\n";
//    cout << t->lpmatch( bytes_4_to_uint32( 209,85,147,99  )  )->str() << "\n";
//    cout << t->lpmatch( bytes_4_to_uint32( 127,0,0,1      )  )->str() << "\n";
//    cout << t->lpmatch( bytes_4_to_uint32( 127,0,0,1      )  )->str() << "\n";
    cout << t->lpmatch( bytes_4_to_uint32( 211,120,0,0  )  )->str() << "\n";
    cout << t->lpmatch( bytes_4_to_uint32( 211,120,0,1  )  )->str() << "\n";


#endif    

    //CITM<AsnPrefix, AsnPrefixLessThan> m;
//    map<int, CITMValue> m;

#if 0
    unordered_map<AsnPrefix, CITMValue> m;

    AsnPrefix p1 = AsnPrefix( 100, Prefix(0xffff0000, 16) );
    m[p1] +=  CITMValue(1,10,100) ;
//    p1 = AsnPrefix( 100, Prefix(0xff000000, 16) );
//    m[p1] +=  CITMValue(1,10,100) ;
//
//    p1 = AsnPrefix( 100, Prefix(0xffff0000, 16) );
//    cout << p1.str() << "," << m[p1].str()   << endl;
//    p1 = AsnPrefix( 100, Prefix(0xff000000, 16) );
//    cout << p1.str() << "," << m[p1].str()   << endl;
//    cout << endl;

#endif

#if 0
//    CITM<AsnPrefix, AsnPrefixLessThan> m;
//    map< AsnPrefix, CITMValue, AsnPrefixLessThan> m;
//
//    AsnPrefix p1 = AsnPrefix( 100, Prefix(0xffff0000, 16) );
//    m.add( p1 ,  CITMValue(1,10,100) );
//    p1 = AsnPrefix( 100, Prefix(0xff000000, 16) );
//    m.add( p1 ,  CITMValue(1,10,100) );
//
//    p1 = AsnPrefix( 100, Prefix(0xffff0000, 16) );
//    cout << p1.str() << "," << m.get(p1).str()   << endl;
//    p1 = AsnPrefix( 100, Prefix(0xff000000, 16) );
//    cout << p1.str() << "," << m.get(p1).str()   << endl;


    typedef map< AsnPrefix, CITMValue, AsnPrefixLessThan> ITMMap;


    ITMMap m;

    AsnPrefix p1 = AsnPrefix( 100, Prefix(0xffff0000, 16) );
    m[p1] +=  CITMValue(1,10,100) ;
    p1 = AsnPrefix( 100, Prefix(0xff000000, 16) );
    m[p1] +=  CITMValue(1,10,100) ;

    p1 = AsnPrefix( 100, Prefix(0xffff0000, 16) );
    cout << p1.str() << "," << m[p1].str()   << endl;
    p1 = AsnPrefix( 100, Prefix(0xff000000, 16) );
    cout << p1.str() << "," << m[p1].str()   << endl;
    cout << endl;

    typeof<>::iterator it;
    for( it = m.begin(); it != m.end(); it++ ) {
        AsnPrefix p =  (*it).first;
        cout <<  p.str() << "," << (*it).second.str() << endl;
    }


#endif

#if 0 // test reading flows from file 
    string fname = "../test/flows/flows1306827770.000000";
    uint32_t counter = 0;

    CFlowSmart flow;
    flow.open( fname );
    if ( ! flow.is_ok() ) error("Could not open file" << fname);
    while( !flow.is_end() ) {   
        if ( flow.next() ) { 
//                cout << " read src " << inet_ntoa(flow.flowbin.ip_src) << endl; // DEBUG
//                cout << " read dst " << inet_ntoa(flow.flowbin.ip_dst) << endl; // DEBUG

            counter++;
        flow.debug();
        cout << counter << " | src " << inet_ntoa(flow.getAddrSrc()) << " :" << flow.getPortSrc() << endl;
        cout << counter << " | dst " << inet_ntoa(flow.getAddrDst()) << " :" << flow.getPortDst() << "  " 
         << flow.getFlows() << " "
         << flow.getOctets() << " "
         << flow.getPackets()  << endl;
         cout << endl;
        }
    }
    flow.close();
#endif



#if 1 // MAIN
    /* parse options */

    int opt, longIndex;
    CTrieTree * trie=NULL;
//    CTrieNode *node1=NULL, *node2=NULL;
    CFlow * flow;


    unordered_map<FlowMapKey*, CITMValue, FlowMapKeyHash, FlowMapKeyEqual> itm;
    unordered_map<FlowMapKey*, CITMValue, FlowMapKeyHash, FlowMapKeyEqual>::iterator itm_it;

#if 0
    AsnPrefix ap  = AsnPrefix( 100, Prefix( 0xffffffff, 16 )  );
    CITMValue val = CITMValue(1,100,10);
    cout << "size " << itm.size() << endl;
    itm[ap] = val;
    cout << "size " << itm.size() << endl;
    cout << ap.str() << " " <<  val.str() << "\n    " <<  itm[ap].str() << endl;
    return 0;
#endif 

    opt = getopt_long( argc, argv, shortOpts, longOpts, &longIndex );
    while ( opt != -1 ) {
        switch(opt) {
            case 'i':
                globalArgs.file_in = optarg;
                break;
            case 'm':
                globalArgs.files_mapping.push_back( string(optarg) );
                break;
            case 'f':
                globalArgs.format = atoi(optarg);
                if ((globalArgs.format < 1) && (globalArgs.format > 2)) { error("--format?"); }
                break;
            case 'k':
                if      ( string(optarg) == "aa" ) globalArgs.keyFormat = KEYFORMAT_ASNASN;
                else if ( string(optarg) == "ap" ) globalArgs.keyFormat = KEYFORMAT_ASNPREFIX;
                else error("--aggregation ?");
                break;
            case 'I':
                if      ( string(optarg) == "nf"      ) globalArgs.inType = INTYPE_NETFLOW;
                else if ( string(optarg) == "netflow" ) globalArgs.inType = INTYPE_NETFLOW;
                else if ( string(optarg) == "s"       ) globalArgs.inType = INTYPE_SMARTBIN;
                else if ( string(optarg) == "smart"   ) globalArgs.inType = INTYPE_SMARTBIN;
                else error("--in-type ?");
                break;

            case '?':
            default:
                display_usage();
                return 0;
                break;
        }
        opt = getopt_long( argc, argv, shortOpts, longOpts, &longIndex );
    }

    display_information();

    if ( ! globalArgs.files_mapping.empty() ) {
     /* load mapping for all the files */
        list<string>::iterator it;
        for( it = globalArgs.files_mapping.begin(); it != globalArgs.files_mapping.end(); it++) {
            /* do for each mapping file */
            if ( trie == NULL )  {
                trie = CTrieTree::newFromFile( *it, CTrieTree::FORMATTXT );
                if ( trie == NULL ) error("Could not load trie " << *it);
            }
            else {   /* append new entries */
                if ( trie->addFromFile( *it, CTrieTree::FORMATTXT ) != CTrieTree::ERR_OK ) {
                    error("Problem with trie...");
                }
            }            
        }
        if ( trie == NULL ) { error("--mapping ? Did not load anything"); }
        cout << "# Trie: nodes all " << trie->getStatsNodesAll() << " set " << trie->getStatsNodesSet() << endl;
        trie->addSpecial();
    }

    if ( trie == NULL && globalArgs.inType == INTYPE_SMARTBIN) { error("--mapping ? Smart format neets AS mapping"); } 

    if ( "" == globalArgs.file_in ) { error("--in ?"); }
    
    /* create ITM map */

    /* open flows */
    switch( globalArgs.inType ) {
        case INTYPE_SMARTBIN: flow = new CFlowSmart(); break;
        case INTYPE_NETFLOW:  flow = new CFlowNF(); break;
        default: assert(0);
    }

    flow->open( globalArgs.file_in );
    if ( ! flow->is_ok() ) error("Could not open" <<  globalArgs.file_in );
    while( !flow->is_end() ) {
        if ( flow->next() ) {   
#if 0
            uint32_t t = flow->getAddrSrc().s_addr;
            cout << endl;
            cout << BYTE3(t) << " " 
                 << BYTE2(t) << " " 
                 << BYTE1(t) << " " 
                 << BYTE0(t) << endl;

            node1 = trie->lpmatch( flow->getAddrSrc() );
    	    cout << inet_ntoa( flow->getAddrSrc() ) << " " << node1->str() << endl;
            node2 = trie->lpmatch( flow->getAddrDst() );
    	    cout << inet_ntoa(     flow->getAddrDst() ) << " " << node2->str() << endl;

          AsnPrefix ap = AsnPrefix( node1->getAsn(), node2->getPrefix() );
            AsnPrefix::hash hash;
            cout << ap.str() << " " << hash(ap) << endl;

#else
            /* for each flow, accumulate it in ITM */
            /* asn from source address */

            uint32_t asn_src = 0, asn_dst = 0;
            Prefix prefix_src, prefix_dst;

            if ( trie ) {
                CTrieNode *node1 = trie->lpmatch( flow->getAddrSrc() ); /* source asn */
                CTrieNode *node2 = trie->lpmatch( flow->getAddrDst() ); /* destination prefix */
                asn_src = node1->getAsn();
                prefix_src = node1->getPrefix();
                asn_dst = node2->getAsn();
                prefix_dst = node2->getPrefix();

            }
            else {
                CFlowNF * nf = dynamic_cast<CFlowNF*>( flow );
                assert(nf);
                asn_src = nf->getAsnSrc();
                prefix_src = Prefix();
                asn_dst = nf->getAsnDst();
                prefix_dst = Prefix();

            }

            FlowMapKey *p_key;
//            AsnPrefix *p_ap;
//            AsnAsn    *p_aa;
//                    p_ap = new AsnPrefix( node1->getAsn(), node2->getPrefix() );
            switch ( globalArgs.keyFormat ) {
//                case KEYFORMAT_ASNPREFIX: p_key = new AsnPrefix( node1->getAsn(), node2->getPrefix() ); break;
//                case KEYFORMAT_ASNASN:    p_key = new AsnAsn(    node1->getAsn(), node2->getAsn() );    break;
                case KEYFORMAT_ASNPREFIX: p_key = new AsnPrefix( asn_src, prefix_dst ); break;
                case KEYFORMAT_ASNASN:    p_key = new AsnAsn(    asn_src, asn_dst    );    break;
                default: assert(0);
            }

            CITMValue val = CITMValue( flow->getFlows(), flow->getOctets(), flow->getPackets() );

// DEBUG 
//            if ( node2->getAsn() == 10056 ) { // DEBUG
//            if ( 1||node1->getAsn() == 4766 ) { // DEBUG
//                CTrieNode *node=NULL;
//
//                uint32_t ip =     node1->getPrefix().getIP();
//                uint8_t netmask = node1->getPrefix().getNetmask();
//
//                node = trie->lpmatch( ip, netmask);
//
//                cout << "DEBUG  " << inet_ntoa( flow->getAddrSrc() ) << " " << node1->str() << " : " << node->str() << endl;
//            }

//            int f_debug = 1;
//            if ( 0 == node2->getPrefix().str().compare("84.88.0.0/25") && node1->getAsn() == 13041010) f_debug = 1;
//            if ( 0 == node2->getPrefix().str().compare("0.0.0.0/0") && node1->getAsn() == 0) f_debug = 1;
//            if ( f_debug) cout << "DEBUG " << ap.str() << " " << val.str() << endl;
//            if ( f_debug ) {
//    	        cout << "SRC " << inet_ntoa( flow->getAddrSrc() ) << " " << node1->str() << endl;
//        	    cout << "DST " << inet_ntoa( flow->getAddrDst() ) << " " << node2->str() << endl;
//            }

//            AsnPrefix::hash hash;


//            itm_it = itm.find(ap);
//            if ( itm_it == itm.end() ) { /* key not found in hash */
//                if (f_debug) cout << "NEW " << hash(ap) << endl;
//                if (f_debug) cout << itm[ap].str() << endl;
//                itm[ap] = val;
//                if (f_debug) cout << itm[ap].str() << endl;
//            } else { /* key exists, accumulate */
//                if (f_debug) cout << "OLD" << endl;
//                itm[ap] += val; // TODO check if accumulate for sure
//                if (f_debug) cout << itm[ap].str() << endl;
//            }


           
//            itm[ ap ] += val; /* is automatically created if not exists */

            itm_it = itm.find( p_key );
            if (itm_it == itm.end() ) { /* not exist -> insert */

                switch ( globalArgs.keyFormat ) {
                    case KEYFORMAT_ASNPREFIX: itm.insert( std::pair<AsnPrefix*,CITMValue>(dynamic_cast<AsnPrefix*>(p_key), val)); break;
                    case KEYFORMAT_ASNASN:    itm.insert( std::pair<AsnAsn*,   CITMValue>(dynamic_cast<AsnAsn*>(p_key),    val)); break;
                    default: assert(0);
                }


//                itm.insert( std::pair<AsnPrefix*,CITMValue>(p_ap, val));
            }
            else { /* add to existing */
                (*itm_it).second += val;
                delete p_key;
            }

#endif

        }
    }
    flow->close();

#if 1
    /* iterate and display */
    if ( globalArgs.keyFormat == KEYFORMAT_ASNPREFIX ) {
        for( itm_it = itm.begin(); itm_it != itm.end(); itm_it++ ) {
            CTrieNode *node;
            AsnPrefix * p_ap = dynamic_cast<AsnPrefix*> ( (*itm_it).first );
            /* match second ASN */
            switch(globalArgs.format) {
                case 1:
                    cout <<  p_ap->getAsn() << " " <<
                             p_ap->getPrefix().str() << " " <<
                             (*itm_it).second.str(" ") << endl;
                    break;
                case 2:
                    node = trie->lpmatch( p_ap->getPrefix().getIP(), p_ap->getPrefix().getNetmask());
                    cout <<  p_ap->getAsn() << " " <<
                             node->getAsn() << " " <<
                             p_ap->getPrefix().str() << " " <<
                             (*itm_it).second.str(" ") << endl;
                    break;
            }
        }
    }
    else if ( globalArgs.keyFormat == KEYFORMAT_ASNASN ) {
        for( itm_it = itm.begin(); itm_it != itm.end(); itm_it++ ) {
            AsnAsn * p_aa = dynamic_cast<AsnAsn*> ( (*itm_it).first );
            cout <<  p_aa->getAsnSrc() << " " <<
                     p_aa->getAsnDst() << " " <<
                     (*itm_it).second.str(" ") << endl;
        }
    }
    else assert(0);
    
#endif

#endif
}

/*
 * TODO add verbose, print ip on unknown match
 * TODO make abstractions of ITM key to fir both prefix and asn

 * TODO read netflows
 * TODO file has BOTH directions, check it with original code


 13041010,84.88.0.0/25,1,57,1
 13041010,84.88.0.0/25,1,61,1
 13041010,84.88.0.0/25,1,63,1

 */

